'''
########################
#
# Copyright 2011, Vanderbilt University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS IS"
# BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
# or implied.  See the License for the specific languaged governing 
# permissions and limitations under the License.
#
########################

@author:Kathy Benitez

updated:July 1, 2011
'''
from sqlalchemy import create_engine, MetaData, Table, Column,\
     String, ForeignKey, Integer, select, func

"""
These functions set up connections to the relevant databases, 
including the foreign key connection necessary to access the values.

connectToDatabaseCounty- The database as indexed by county is smaller
 than the Zip database.  Therefore, we use the county database not only
 when we want access to county information, but also when we don't want
 or need ZIP code information.

connectToDatabaseZip

"""
class PopStatsConnection():
    def __init__(self, engine="", uname="", passwd="", hostAndPort="", database="", echo=False, url=""):
        if url:
            sqlengine = create_engine(url)
        else:
            sqlengine = create_engine('%s://%s:%s@%s/%s'%\
                                  (engine, uname, 
                                   passwd, hostAndPort, 
                                   database), echo=echo)
        self.metadata = MetaData()
        self.xn = sqlengine.connect()
        self.summaryFields = None
        self.records = None
        self.description = Table('fieldnames', self.metadata,
                           Column('id', Integer, primary_key=True),
                           Column('fieldname', String(10), unique=True, index=True),
                           #Column('race', Enum('American Indian and Alaska Native alone','Asian alone','Black or African American alone','Native Hawaiian and Other Pacific Islander alone','Some other race alone','Two or more races','White alone')),
                           Column('race', String(50)),
                           #Column('sex', Enum('Female','Male')),
                           Column('sex', String(6)),
                           Column('minage', Integer),
                           Column('maxage', Integer))   
    
    def tearDown(self):
        self.xn.close()
        self.sqlengine.close()
    """
    Input: A dictionary of category: possible values.
    If category is not in the dictionary, it is desired with full
    generality.

    """
    def setupQuery(self):
        s = select([self.description.c.minage,
                    self.description.c.maxage,
                    self.description.c.race,
                    self.description.c.sex,
                    func.sum(self.records.c.population).label("pop")],
                   from_obj=[self.records.join(self.description)])
        return s
    def createTables(self):
        self.metadata.create_all()
    def howManyMatches(self, specifications):
        q = self.setupQuery()
        if specifications:
            for field in specifications:
                if field=="age":
                    """ We're going to deal with age separately """
                    continue
                if field in self.description.c:
                    q = q.where(eval("self.description.c."+field).in_(specifications[field]))
                elif field in self.records.c:
                    q = q.where(eval("self.records.c."+field).in_(specifications[field]))
                else:
                    raise TypeError(field + " is not a field covered by the population statistics")
        
            if "age" in specifications:
                q = q.where(self.description.c.maxage >= min(specifications["age"]))
                q = q.where(self.description.c.minage <= max(specifications["age"]))
                q = q.group_by(self.description.c.minage, self.description.c.maxage)
                result = self.xn.execute(q)
                matches = []
                for r in result:
                    """ count how many ages specified are in the age range """
                    minage = r[self.description.c.minage]
                    maxage = r[self.description.c.maxage]
                    match = [a <= maxage and a >= minage for a in specifications["age"]].count(True)
                    matches.append((match/(maxage-minage + 1.)) * int(r["pop"]))
                result.close()
                return sum(matches,0)
        
        result = self.xn.execute(q)
        matches = [r["pop"] for r in result]
        result.close()
        return sum(matches, 0)

class PopStatsCounty(PopStatsConnection):
    def __init__(self, engine="", uname="", passwd="", hostAndPort="", database="", echo=False, url=""):
        PopStatsConnection.__init__(self, engine, uname, passwd, hostAndPort, database, echo, url)
        self.records = Table('by_county', self.metadata,
                   Column('id', Integer, primary_key=True),
                   Column('county', String(30)),
                   Column('state', String(30)),
                   Column('fieldname_id', String(10), ForeignKey(self.description.c.fieldname)),
                   Column('population', Integer))       
    def setupQuery(self):
        s = PopStatsConnection.setupQuery(self)
        s.append_column(self.records.c.county)
        s.append_column(self.records.c.state)
        return s   
    def testIntegrity(self):
        popByState2000={"California":33871648,
                        "Texas":20851820,
                        "New York":18976457,
                        "Florida":15982378,
                        "Illinois":12419293,
                        "Washington":5894121,
                        "Tennessee":5689283,
                        "Wisconsin":5363675,
                        "Minnesota":4919479}
        result = select([self.records.c.state,
                         func.sum(self.records.c.population).label("Pop")]).group_by(self.records.c.state).execute()
        matches = [popByState2000[row["state"]] == row["Pop"] for row in result if row["state"] in popByState2000]
        return all(matches) 

class PopStatsZip(PopStatsConnection):
    def __init__(self, engine="", uname="", passwd="", hostAndPort="", database="", echo=False, url=""):
        PopStatsConnection.__init__(self, engine, uname, passwd, hostAndPort, database, echo, url) 
        self.records = Table('by_zip', self.metadata,
                   Column('id', Integer, primary_key=True),
                   Column('zip5', String(5)),
                   Column('fieldname_id', String(10), ForeignKey(self.description.c.fieldname)),
                   Column('population', Integer))
    def setupQuery(self):
        s = PopStatsConnection.setupQuery(self)
        s.append_column(self.records.c.zip5)
        return s   
    def testIntegrity(self):
        # Numbers for the integrity check taken from American Fact Finder, for some zips local to
        # Nashville.
        popByZip2000={37201:1167,
                      37203:12871,
                      37204:11024,
                      37205:21861,
                      37206:27751,
                      37207:35744}
        result = select([self.records.c.zip5,
                         func.sum(self.records.c.population).label("Pop")]).group_by(self.records.c.zip5).execute()
        matches = [popByZip2000[row["zip5"]] == row["Pop"] for row in result if row["zip5"] in popByZip2000]
        return all(matches)